import argparse
import numpy as np
import os
import torch

from model import PointPillarsV2

def main(args):
    CLASSES = {
        'Pedestrian': 0, 
        'Cyclist': 1, 
        'Car': 2
        }
    if not os.path.exists(args.ckpt):
        raise FileNotFoundError
    if not args.no_cuda:
        model = PointPillarsV2(nclasses=len(CLASSES), use_intensity=args.use_intensity).cuda()
        checkpoint = torch.load(args.ckpt)
        model.load_state_dict(checkpoint['state_dict'])
    else:
        model = PointPillarsV2(nclasses=len(CLASSES))
        model.load_state_dict(
            torch.load(args.ckpt, map_location=torch.device('cpu')))

    # export vfe onnx model.
    vfe_model = model.pillar_encoder.pfn_layer
    vfe_model.cuda()
    vfe_model.eval()
    with torch.no_grad():
        pts_feats = torch.zeros((40000, 8, 32), dtype=torch.float32, device='cuda:0')
        torch.onnx.export(vfe_model, pts_feats, "vfe.onnx", opset_version=11,
                          input_names=['vfe_input'], output_names=['vfe_output'])
    
    # export rpn backbone net.
    rpn_backbone = model.rpn_backbone
    rpn_backbone.cuda()
    rpn_backbone.eval()
    with torch.no_grad():
        pillars_feats = torch.zeros((1, 64, 480, 240), dtype=torch.float32, device='cuda:0')
        torch.onnx.export(rpn_backbone, pillars_feats, "rpn_backbone.onnx", opset_version=11,
                          input_names=['rpn_input'], output_names=['cls_pred', 'reg_pred', 'dir_pred'])


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Configuration Parameters')
    parser.add_argument('--ckpt', default='pretrained/xxx.pth', help='your checkpoint for kitti')
    parser.add_argument('--use_intensity', type=bool, default=False)
    parser.add_argument('--no_cuda', action='store_true',
                        help='whether to use cuda')
    args = parser.parse_args()

    main(args)